import os
import time
import numpy as np
from scipy.io import loadmat
from PIL import Image
from tqdm import tqdm

from mit_utils import cfg
from mit_utils import ValDataset, DataLoader
from mit_utils import AverageMeter, colorEncode, accuracy, intersectionAndUnion, setup_logger
from mit_utils import as_numpy

import jittor as jt

from model import EncoderDecoder


colors = loadmat('data/color150.mat')['colors']

def visualize_result(data, pred, dir_result):
    (img, seg, info) = data

    # segmentation
    seg_color = colorEncode(seg, colors)

    # prediction
    pred_color = colorEncode(pred, colors)

    # aggregate images and save
    im_vis = np.concatenate((img, seg_color, pred_color),
                            axis=1).astype(np.uint8)

    img_name = info.split('/')[-1]
    Image.fromarray(im_vis).save(
        os.path.join(dir_result, img_name.replace('.jpg', '.png')))

@jt.no_grad()
@jt.single_process_scope()
def evaluate(segmentation_module, loader, cfg):
    acc_meter = AverageMeter()
    intersection_meter = AverageMeter()
    union_meter = AverageMeter()
    time_meter = AverageMeter()

    segmentation_module.eval()

    pbar = tqdm(total=len(loader))
    for batch_data in loader:
        # process data
        batch_data = batch_data[0]
        seg_label = as_numpy(batch_data['seg_label'][0])
        img_resized_list = batch_data['img_data']

        tic = time.perf_counter()
        segSize = (seg_label.shape[0], seg_label.shape[1])
        scores = jt.zeros(1, cfg.DATASET.num_class, segSize[0],
                                segSize[1])

        for img in img_resized_list:
            feed_dict = {
                'img_data': img,
                'seg_label': batch_data['seg_label']
            }
            scores_tmp = segmentation_module(feed_dict, segSize=segSize)
            scores = scores + scores_tmp / len(cfg.DATASET.imgSizes)

        pred = jt.argmax(scores, dim=1)[0]
        pred = pred.squeeze(0).clone().numpy()

        time_meter.update(time.perf_counter() - tic)

        # calculate accuracy
        acc, pix = accuracy(pred, seg_label)
        intersection, union = intersectionAndUnion(pred, seg_label,
                                                   cfg.DATASET.num_class)
        acc_meter.update(acc, pix)
        intersection_meter.update(intersection)
        union_meter.update(union)

        # visualization
        if cfg.VAL.visualize:
            visualize_result(
                (batch_data['img_ori'], seg_label, batch_data['info']), pred,
                os.path.join(cfg.DIR, 'result'))

        pbar.update(1)

    # summary
    iou = intersection_meter.sum / (union_meter.sum + 1e-10)
    for i, _iou in enumerate(iou):
        print('class [{}], IoU: {:.4f}'.format(i, _iou))

    print('[Eval Summary]:')
    print(
        'Mean IoU: {:.4f}, Accuracy: {:.2f}%, Inference Time: {:.4f}s'.format(
            iou.mean(),
            acc_meter.average() * 100, time_meter.average()))


def main(cfg):
    # TODO : CHANGE TO YOU CHECKPOINT PATH
    resume = "/data/share/leixy/ccnet_jittor/ckpt/ade20k-resnet101-new-cca_deepsup/epoch_27.pkl"
    segmentation_module = EncoderDecoder(resume=resume)

    # Dataset and Loader
    dataset_val = ValDataset(cfg.DATASET.root_dataset, cfg.DATASET.list_val,
                             cfg.DATASET)
    loader_val = DataLoader(dataset_val,
                            batch_size=1,
                            shuffle=False,
                            collate_fn=lambda x: x,
                            num_workers=cfg.TRAIN.workers,
                            drop_last=True)

    # Main loop
    print('Begin evaluation!')
    
    evaluate(segmentation_module, loader_val, cfg)

    print('Evaluation Done!')


if __name__ == '__main__':
    jt.flags.use_cuda = 1
    jt.flags.device_id = 6
    
    cfg.merge_from_file("config/eval_ade20k-resnet101dilated-cca_deepsup.yaml") # this config file is different from that for training
    
    logger = setup_logger(distributed_rank=0)

    if not os.path.isdir(os.path.join(cfg.DIR, "result")):
        os.makedirs(os.path.join(cfg.DIR, "result"))

    main(cfg)
